import re

import chromadb
from chromadb.utils import embedding_functions
from docx import Document
import os
from openai import OpenAI


def load_employee_file(file_path):
    document = Document(file_path)
    all_text = []
    for para in document.paragraphs:
        clean_text = para.text.strip()
        clean_text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s\.,;:!?()""''\\-]', '', clean_text)
        clean_text = re.sub(r'\s+', ' ', clean_text)
        all_text.append(clean_text)
    return '\n'.join(all_text)


def split_text(text, chunk_size=200):
    chunks = []
    for i in range(0, len(text), chunk_size):
        chunks.append(text[i:i + chunk_size])
    return chunks


def text_embeddings(texts):
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),  # 如果您没有配置环境变量，请在此处用您的API Key进行替换
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"  # 百炼服务的base_url
    )

    # 如果输入是字符串，转换为列表
    if isinstance(texts, str):
        texts = [texts]

    # 阿里云百炼服务的text-embedding-v4模型批处理大小不能超过10
    batch_size = 10
    embeddings = []

    # 分批处理输入文本
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        completion = client.embeddings.create(
            model="text-embedding-v4",
            input=batch,
            dimensions=1024,  # 指定向量维度（仅 text-embedding-v3及 text-embedding-v4支持该参数）
            encoding_format="float"
        )

        print(f"Processing batch {i + 1}:{completion}")

        # 提取嵌入向量
        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


def store_in_chromadb(chunks, collection_name, persist_directory="./chroma_db"):
    # 1. 生成文本嵌入向量
    embeddings = text_embeddings(chunks)
    if len(embeddings) != len(chunks):
        print("嵌入向量生成异常，数量与文本块不匹配")
        return

    # 2. 初始化ChromaDB客户端并创建/获取集合
    client = chromadb.PersistentClient(path=persist_directory)
    # 创建集合时指定嵌入函数（此处使用已生成的嵌入，故设为None）
    collection = client.get_or_create_collection(
        name=collection_name,
        embedding_function=None  # 不使用内置嵌入函数，使用自定义生成的嵌入
    )

    # 3. 生成唯一ID（格式：chunk_0, chunk_1, ...）
    ids = [f"chunk_{i}" for i in range(len(chunks))]

    # 4. 将数据存入集合
    collection.add(
        documents=chunks,  # 文本片段
        embeddings=embeddings,  # 对应的嵌入向量
        ids=ids  # 唯一标识
    )

    print(f"成功存入ChromaDB集合 '{collection_name}'，共 {len(chunks)} 个文本块")


def query_chroma(query_text):
    # 创建嵌入函数
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        model_name="text-embedding-v4",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    # 2. 初始化ChromaDB客户端并创建/获取集合
    persist_directory = "./chroma_db"
    chroma_client = chromadb.PersistentClient(path=persist_directory)
    # 获取集合时指定嵌入函数
    collection = chroma_client.get_collection(
        name="employee_collection",
        embedding_function=embedding_function
    )

    # 查询与特定文本相似的向量
    results = collection.query(
        query_texts=[query_text],
        n_results=5,  # 返回最相似的5个结果
        include=['embeddings', 'documents', 'metadatas', 'distances']  # 显式指定需要返回的字段
    )
    # {'ids': [['text_41', 'text_33', 'text_37', 'text_35', 'text_44']], 'embeddings': None, 'documents': [['年假（5-15天，按工作年限）、病假（按国家规定）、婚假/产假/陪产假等法定假期\n健康福利：每年1次', '到岗\n相应调整下班时间，保证8小时工作时长\n需要提前向直属主管申请并获得批准\n特殊工作安排\n项目紧急', '分钟以上\n早退：提前离开工作岗位\n旷工：未经批准不到岗工作\n考勤处理\n迟到3次以内：口头提醒\n迟到3', '要求\n员工应按时上下班，严格遵守考勤制度\n使用公司考勤系统进行打卡签到\n因故不能按时到岗应提前请假\n', '、年度旅游/聚餐、员工家庭日、兴趣社团支持\n第六章 培训发展\n6.1 培训体系\n新员工培训（3周）\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[None, None, None, None, None]], 'distances': [[0.6660666465759277, 0.8344534039497375, 0.9031039476394653, 0.9283781051635742, 0.9392238855361938]]}

    print(results)
    return results


def rag_answer_question(results, question):
    retrieved_chunks = list(set(results["documents"][0]))
    print("retrieved_chunks", retrieved_chunks)

    # 构建上下文（把片段拼起来，告诉LLM"基于这些内容回答"）
    context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
    print(f"\n【检索到的参考内容】\n{context}\n")
    prompt = f"""
    你是一个专业的HR助理，请仔细阅读以下员工手册内容并准确回答用户问题。

    【员工手册内容】：
    {context}

    【用户问题】：
    {question}

    【回答要求】：
    1. 严格基于员工手册内容回答，不要编造或推测
    2. 如果手册中没有相关信息，请明确说明"手册中未提及此信息"
    3. 回答应简洁明了，条理清晰
    4. 如有多个要点，请分点列出
    5. 如果涉及具体条款，可引用相关内容
    """

    print("提示词:", prompt)

    messages = [{"role": "user", "content": prompt}]
    ai_response = get_response(messages)
    print("大模型返回的答案:", ai_response)


def get_response(messages):
    client = OpenAI(
        # 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：api_key="sk-xxx",
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        # base_url="http://localhost:11434/v1",

    )

    completion = client.chat.completions.create(
        # model="qwen:latest",
        model="qwen-plus",
        messages=messages,
    )
    return completion.choices[0].message.content


if __name__ == '__main__':
    # employee_file_text = load_employee_file("employee_manual.docx")
    # print(employee_file_text)
    # chunks = split_text(employee_file_text)
    # print(chunks)
    # store_in_chromadb(chunks, "employee_collection")
    question = "员工工作满3年，年假有几天？需要提前多久申请"
    results = query_chroma(question)
    rag_answer_question(results, question)
